local ls = require("luasnip")
local s = ls.snippet
local t = ls.text_node
local i = ls.insert_node
local fmt = require("luasnip.extras.fmt").fmt

ls.add_snippets("python", {
    -- Print
    s("st", fmt("print({})", {i(1)})),

    -- 函数定义
    s("def", fmt([[
def {}({}) -> {}:
    {}
    ]], {
        i(1, "function_name"),
        i(2, "params"),
        i(3, "None"),
        i(4, "pass"),
    })),

    -- 类定义
    s("class", fmt([[
class {}:
    def __init__(self, {}):
        {}
        
    {}
    ]], {
        i(1, "ClassName"),
        i(2, "params"),
        i(3, "pass"),
        i(4, "# methods"),
    })),

    -- if __name__ == "__main__"
    s("ifmain", fmt([[
if __name__ == "__main__":
    {}
    ]], {
        i(1, "main()"),
    })),

    -- try-except
    s("try", fmt([[
try:
    {}
except {} as e:
    {}
    ]], {
        i(1, "# code"),
        i(2, "Exception"),
        i(3, "print(f'Error: {e}')"),
    })),

    -- with open
    s("with", fmt([[
with open({}, {}) as f:
    {}
    ]], {
        i(1, "'file.txt'"),
        i(2, "'r'"),
        i(3, "# code"),
    })),

    -- list comprehension
    s("lc", fmt("[{} for {} in {}]", {
        i(1, "x"),
        i(2, "x"),
        i(3, "iterable"),
    })),

    -- dictionary comprehension
    s("dc", fmt("{{{}: {} for {} in {}}}", {
        i(1, "key"),
        i(2, "value"),
        i(3, "x"),
        i(4, "iterable"),
    })),

    -- async function
    s("adef", fmt([[
async def {}({}) -> {}:
    {}
    ]], {
        i(1, "function_name"),
        i(2, "params"),
        i(3, "None"),
        i(4, "pass"),
    })),

    -- import
    s("imp", fmt("from {} import {}", {
        i(1, "module"),
        i(2, "name"),
    })),
})
